package complexdatastructure.unionfindset;

import java.util.*;

/**
 * @Classname : EvaluateDivision
 * @Description : 399. 除法求值
 * https://leetcode.cn/problems/evaluate-division/
 * @Author : cty
 * @Date 2023/4/12 22:36
 */


public class EvaluateDivision {

    /**
     * 测试
     * @param args
     */
    public static void main(String[] args) {
        List<List<String>> equations = new ArrayList<>();
        equations.add(Arrays.asList("a", "b"));
        equations.add(Arrays.asList("b", "c"));
        double[] values = {2.0, 3.0};
        List<List<String>> queries = new ArrayList<>();
        queries.add(Arrays.asList("a", "c"));
        queries.add(Arrays.asList("b", "a"));
        queries.add(Arrays.asList("a", "e"));
        queries.add(Arrays.asList("a", "a"));
        queries.add(Arrays.asList("x", "x"));

        EvaluateDivision evaluateDivision = new EvaluateDivision();
        double[] res = evaluateDivision.calcEquation(equations, values, queries);
        System.out.println(res);  // res: [6.0, 0.5, -1.0, 1.0, -1.0]
    }


    Map<String, Integer> indexMap;  // 记录字符串编号
    Map<String, Double> valuesMap;  // 记录问题答案

    public double[] calcEquation(List<List<String>> equations, double[] values, List<List<String>> queries) {
        // 初始化两个记录Map
        indexMap = new HashMap<>();
        valuesMap = new HashMap<>();
        for (int i = 0; i < equations.size(); i++) {
            List<String> equation = equations.get(i);
            String x = equation.get(0), y = equation.get(1);
            indexMap.putIfAbsent(x, indexMap.size());
            indexMap.putIfAbsent(y, indexMap.size());
            valuesMap.putIfAbsent(indexMap.get(x)+"#"+indexMap.get(y), values[i]);
            valuesMap.putIfAbsent(indexMap.get(y)+"#"+indexMap.get(x), 1.0 / values[i]);
            valuesMap.putIfAbsent(indexMap.get(x)+"#"+indexMap.get(x), 1.0 / values[i]);
            valuesMap.putIfAbsent(indexMap.get(y)+"#"+indexMap.get(y), 1.0 / values[i]);
        }
        // 初始化并查集
        UF uf = new UF(indexMap.size());
        for (List<String> equation : equations) {
            String x = equation.get(0), y = equation.get(1);
            uf.union(indexMap.get(x), indexMap.get(y));
        }
        // 遍历元素，压缩并查集层级
        for (int v : indexMap.values()) {
            uf.find(v);
        }
        // 封装返回结果
        double[] res = new double[queries.size()];
        for (int i = 0; i < queries.size(); i++) {
            List<String> query = queries.get(i);
            Integer x = indexMap.get(query.get(0)), y = indexMap.get(query.get(1));
            // 若至少有一个字符串之前未出现 或 两个字符串不连通，则无法确定答案
            if ((x == null || y == null)
                    || (uf.find(x) != uf.find(y))) {
                res[i] = -1.0;
            } else {
                Double value = valuesMap.get(x+"#"+y);
                Integer p = uf.find(x);
                if (x.equals(y)) res[i] = 1.0;  // 两个字符串相同
                else if (value != null) res[i] = value;  // 两个字符串在知识数组出现过
                else if (x.equals(p)) res[i] = value;  // 左字符串为二者共同父节点
                else if (y.equals(p)) res[i] = 1.0 / valuesMap.get(y+"#"+x);  // 右字符串为二者共同父节点
                else res[i] = valuesMap.get(p+"#"+y) / valuesMap.get(p+"#"+x);  // 一般情况
                valuesMap.putIfAbsent(x+"#"+y, res[i]);  // 积累知识
            }
        }
        return res;
    }

    class UF {
        private int[] parent;
        private int[] rank;

        public UF(int n) {
            parent = new int[n];
            rank = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                rank[i] = 1;
            }
        }

        public int find(int x) {
            if (parent[x] == x) {
                return x;
            } else {
                int pp = find(parent[x]);
                int p = parent[x];
                parent[x] = pp;
                setValue(pp, p, x);
                return pp;
            }
        }

        public boolean union(int x, int y) {
            int xp = find(x);
            int yp = find(y);
            if (xp == yp) return false;
            if (rank[xp] >= rank[yp]) {
                if (rank[xp] == rank[yp]) rank[xp]++;
                parent[y] = xp;
                setValue(xp, x, y);
                parent[yp] = xp;
                setValue(xp, y, yp);

            } else {
                parent[x] = yp;
                setValue(yp, y, x);
                parent[xp] = yp;
                setValue(yp, x, xp);
            }
            return true;
        }
    }

    // 计算b合并到ap之下的答案
    private void setValue(int ap, int a, int b) {
        if (valuesMap.get(ap+"#"+b) != null || ap == a) return;
        double apaValue = valuesMap.get(ap+"#"+a);
        double abValue = valuesMap.get(a+"#"+b);
        double apbValue = apaValue * abValue;
        valuesMap.putIfAbsent(ap+"#"+b, apbValue);
        valuesMap.putIfAbsent(b+"#"+ap, 1.0 / apbValue);
    }
}
